import numpy as np
import cv2
import caliou
res10_300x300_ssd_iter_140000_fp16 = cv2.dnn.readNetFromCaffe("deploy.prototxt","res10_300x300_ssd_iter_140000_fp16.caffemodel")

# face count
def detect_(image, conf_threshold=0.5):
	pad = 90
	image = cv2.resize(image, (120, 120))
	image = cv2.copyMakeBorder(image,pad,pad,pad,pad,cv2.BORDER_CONSTANT,value=[0,0,0])
	(h, w) = image.shape[:2]
	blob = cv2.dnn.blobFromImage(cv2.resize(image, (300, 300)), 1.0,
		(300, 300), (104.0, 177.0, 123.0))

	res10_300x300_ssd_iter_140000_fp16.setInput(blob)
	detections = res10_300x300_ssd_iter_140000_fp16.forward()
	bboxes = []
	for i in range(detections.shape[2]):
		confidence = detections[0, 0, i, 2]
		if confidence < conf_threshold:
			continue
		box = detections[0, 0, i, 3:7] * np.array([w, h, w, h])
		(startX, startY, endX, endY) = box.astype("int")
		bboxes.append([startX, startY, endX, endY])
		text = "{:.2f}%".format(confidence * 100)
		y = startY - 10 if startY - 10 > 10 else startY + 10
		cv2.rectangle(image, (startX, startY), (endX, endY),
			(0, 0, 255), 2)
		cv2.putText(image, text, (startX, y),
			cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 0, 255), 2)

	print(len(bboxes))
	print(bboxes)
	bboxes = caliou.calculateIoU(bboxes)
	cv2.imshow("Output", image)
	cv2.waitKey(0)

detect_(cv2.imread('old.jpg'))
detect_(cv2.imread('o1.jpg'))
detect_(cv2.imread('double.jpg'))
detect_(cv2.imread('double1.jpg'))
